import argparse
import os
from utils.angr_util import AngrUtil


def run(args):
    path = args.path
    if not os.path.exists(path):
        raise Exception("File Not Found!")

    if os.path.isdir(path):
        raise Exception("Is A Dir!")

    util = AngrUtil(path)
    util.identify_function()
    print(util.function_names)

    idx = args.func_idx
    if idx >= len(util.function_names):
        print("Index is too large!")
        idx = 0
    g = util.get_function_graph(util.function_names[idx])
    if args.extend == 1:
        blocks = util.get_func_blocks_by_name(util.function_names[idx])
        blocks_dict = util.contruct_block_dict(blocks)
        g, blocks_dict = util.try_extend_graph(g, blocks_dict)
        print(blocks_dict.keys())
        # trace = util.random_walk(blocks_dict, g)

    util.plot_function_graph(g, "cfg")


if __name__ == "__main__":
    parser = argparse.ArgumentParser("Plot The CFG")
    parser.add_argument("-p", "--path", type=str, help="The path of .so data")
    parser.add_argument("-e",
                        "--extend",
                        type=int,
                        default=0,
                        help="The trigger of extending the graph")
    parser.add_argument("-f",
                        "--func_idx",
                        type=int,
                        default=0,
                        help="The index of compute functions")
    args = parser.parse_args()
    run(args)
